package main

import "log"

func main() {
	arr := []int{3, 2, 1, 5, 6, 4}
	k := 2
	topK := TopKPartition(arr, k)
	log.Println(topK)
}

// 寻找最大的K个数
func partition(arr []int, k int, result []int) []int {
	if len(arr) <= 1 {
		log.Printf("Array length is %d, no need to sort.", len(arr))
		return result
	}

	pivotIndex := 0
	left := 1
	right := len(arr) - 1

	for left < right {
		// 先移动right
		for left < right && arr[right] >= arr[pivotIndex] {
			right--
		}
		// 后移动left
		for left < right && arr[left] <= arr[pivotIndex] {
			left++
		}
		if left < right {
			arr[left], arr[right] = arr[right], arr[left]
		}
	}

	// 如果left == right，则left或right都指向pivot，此时将pivot与left或right交换，即可将pivot放到正确位置
	if arr[left] < arr[pivotIndex] {
		arr[left], arr[pivotIndex] = arr[pivotIndex], arr[left]
		pivotIndex = left
	}

	m := len(arr) - pivotIndex
	n := len(arr) - pivotIndex - 1
	if m == k { // pivot+右边的元素正好有k个元素
		result = append(result, arr[pivotIndex:]...)
	} else if n == k {  // pivot右边有k-1个元素
		result = append(result, arr[pivotIndex+1:]...)
	} else if n > k { // pivot右边多于k个元素
		result = partition(arr[pivotIndex+1:], k, result)
	} else { // pivot右边少于k个元素则继续partition右边都属于，对左边来就行
		result = append(result, arr[pivotIndex:]...)
		result = partition(arr[:pivotIndex], k-m, result)
	}
	return result
}

func TopKPartition(arr []int, k int) []int {
	result := make([]int, 0, k)
	brr := make([]int, len(arr))
	for i, ele := range arr {
		brr[i] = ele
	}
	return partition(brr, k, result)
}